'''
Author: goog
Date: 2021-12-13 16:45:46
LastEditTime: 2022-01-13 09:35:44
LastEditors: goog
Description: 下段检测只在10号摄像头检测
FilePath: /chengdu/TensorRT20220110/DetectionZM/ZST/ZST.py
Time Limit Exceeded!
'''
import json
import os
import cv2
import torch
import numpy as np
from PIL import Image
from torchvision import transforms
from torchvision.models import resnet18
from sklearn.decomposition import PCA


def zst(cfg, image, barcode, classif_net, visulization=False):
    defect = "ZST"
    tmp = cv2.cvtColor(image['side'].copy(), cv2.COLOR_BGR2RGB) 
    

    abs_dir = cfg['abs_dir']
    config_path = os.path.join(abs_dir, 'DetectionZM/'+defect+'/'+defect+'.json')
    with open(config_path, 'r') as fp:
        config = json.load(fp)
        # if barcode not in config.keys():
        #     barcode = 'C50'+barcode[3:]
        if barcode[:3] in ['C50', 'C52', 'C53', 'C54', 'C55', 'C56', 'C57', 'C58']:
            # WGW
            barcode = 'C50'+barcode[3:]
        elif barcode[:3] in ['C51']:
            # FMW
            barcode = 'C51'+barcode[3:]
        elif barcode[:3] in ['C59', 'C60', 'H1H', 'H2H', 'C61']:
            # ZMW
            barcode = 'C59'+barcode[3:]
        elif barcode[:3] in ['H3H']:
            # OMW
            barcode = 'H3H'+barcode[3:]
        else:
            barcode = 'C50'+barcode[3:]


        points = config[barcode]['points']['1']
        points1 = config[barcode]['points']['2']
        points2 = config[barcode]['points']['3']
    
    roi = tmp[points[1]: points[1]+points[3], points[0]: points[0]+points[2], :]
    roi1 = tmp[points1[1]: points1[1]+points1[3], points1[0]: points1[0]+points1[2], :]
    roi2 = tmp[points2[1]: points2[1]+points2[3], points2[0]: points2[0]+points2[2], :]
    
    if visulization:
        cv2.imwrite(os.path.join(cfg['abs_dir'], 'DetectionZM/'+defect+'/'+defect+'.png'), roi)
        cv2.imwrite(os.path.join(cfg['abs_dir'], 'DetectionZM/'+defect+'/'+defect+'1.png'), roi)
        cv2.imwrite(os.path.join(cfg['abs_dir'], 'DetectionZM/'+defect+'/'+defect+'2.png'), roi)

    _, color = classif_net.clsZSTColor(roi, roi1, roi2)
    print(color)
    return color == config[barcode]['class']


            
            

